# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     metric
   Description :   
   Author :       lth
   date：          2022/1/24
-------------------------------------------------
   Change Activity:
                   2022/1/24 12:30: create this script
-------------------------------------------------
this script is used to calculate the performance of the weight
"""
__author__ = 'lth'

from math import exp

import torch
import torch.nn.functional as F

from utils import denormalize


class Metric:
    def __init__(self):
        self.generate_image = None
        self.hr_image = None

        self.window_size = 11

    def update(self, generate_image, hr_image):
        self.generate_image = generate_image.float()
        self.hr_image = hr_image.float()

    def SSIM(self):
        g = self.cover_image(self.generate_image)
        h = self.cover_image(self.hr_image)

        min_val = 0
        max_val = 255
        L = max_val - min_val
        h = torch.clamp(h, 0, max_val)

        padd = 0
        (_, channel, height, width) = g.size()

        real_size = min(self.window_size, height, width)
        window = Metric.create_window(real_size, channel=channel).to(g.device)

        mu1 = F.conv2d(g, window, padding=padd, groups=channel)
        mu2 = F.conv2d(h, window, padding=padd, groups=channel)

        mu1_sq = mu1.pow(2)
        mu2_sq = mu2.pow(2)
        mu1_mu2 = mu1 * mu2

        sigma1_sq = F.conv2d(g * g, window, padding=padd, groups=channel) - mu1_sq
        sigma2_sq = F.conv2d(h * h, window, padding=padd, groups=channel) - mu2_sq
        sigma12 = F.conv2d(g * h, window, padding=padd, groups=channel) - mu1_mu2

        C1 = (0.01 * L) ** 2
        C2 = (0.03 * L) ** 2

        v1 = 2.0 * sigma12 + C2
        v2 = sigma1_sq + sigma2_sq + C2
        cs = torch.mean(v1 / v2)  # contrast sensitivity
        ls = (2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)  # luminance
        ss = (sigma12 + C2 / 2) / (torch.sqrt(sigma2_sq) * torch.sqrt(sigma1_sq) + C2 / 2)  # structure

        ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)

        ssim_map = ssim_map.mean()

        return ssim_map, cs, ls.mean(), ss.mean()

    def PSNR(self):
        g = self.cover_image(self.generate_image)
        h = self.cover_image(self.hr_image)
        max_pixel = 255.0

        h = torch.clamp(h, 0, max_pixel)

        return 10 * self.tf_log10(max_pixel ** 2) / (torch.mean(torch.pow(g - h, 2)))

    @staticmethod
    def tf_log10(x):
        numerator = torch.log(torch.tensor(x))
        denominator = torch.log(torch.tensor(10.0))
        return numerator / denominator

    def cover_image(self, image):
        return denormalize(image, mean=0.5, std=0.5)

    @staticmethod
    def create_window(window_size, channel=1):
        _1D_window = Metric.gaussian(window_size, 1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
        return window

    @staticmethod
    def gaussian(window_size, sigma):
        gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
        return gauss / gauss.sum()
